# -*- coding: utf-8 -*-
"""
Created on Mon Jun 28 23:12:49 2021

@author: zhuo 木鸟
"""

import pandas as pd
import numpy as np
from statsmodels.tsa.arima_model import ARIMA
from series_to_supervised import series_to_supervised
from autoregression_analyze import plot_result
from moving_average_analyze import plot_residual_error
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
from calculate_plate_index import plot_plate_index

plt.rcParams['font.sans-serif']=['SimHei']    #画图时使用中文字体
plt.rcParams['axes.unicode_minus'] = False

def plot_residual_kde(residual):
    
    fig = plt.figure()
    residual.plot.kde(legend=False)
    # plt.xlim(left=-30, right=30)
    plt.grid()
    plt.xlabel('残差', fontsize=16)
    plt.ylabel('概率密度', fontsize=16)
    plt.title('核密度图', fontsize=16)
    fig.savefig('../图片/ARIMA的残差概率密度图.png')
    plt.show()


if __name__ == '__main__':
    path = r'../附件/中间数据/光伏-建筑板块市值.xlsx'
    series = pd.read_excel(path, header=0, index_col=0, parse_dates=True, squeeze=True)
    
    train, test = series['2019-4-1':'2021-4-30'], series['2021-5-1':'2021-5-28']
    # p, d, q
    model = ARIMA(train, order=(2,1,0)).fit(disp=0)
    print('模型的参数如下： \n', model.summary())
    
    train_resid = pd.DataFrame(model.resid)
    # 画出误差·
    plot_residual_error(train_resid.index, train_resid.values)
    
    # 画出误差的核密度图
    plot_residual_kde(train_resid)
    
    y_train_pred = model.predict(start=1, end=760, typ='levels')
    print('模型在训练集中的 MSE 为: ', mean_squared_error(train.iloc[1:], y_train_pred))
    
    y_test_pred = model.predict(start=761, end=787, dynamic=True, typ='levels')
    print('模型在测试集中的 MSE 为: ', mean_squared_error(test, y_test_pred))
    
    # 第二问结果
    model = ARIMA(train, order=(2,0,0)).fit(disp=0)
    y_pred = model.predict(start=761, end=761+365, dynamic=True, typ='levels')
    
    plot_plate_index(y_pred, future=True)
    for window in [20,21,60]:
        plate_index_MA = y_pred.rolling(window=window).mean()
        plot_plate_index(plate_index_MA, window=window, future=True)